use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{parse_macro_input, ItemFn, Meta, Lit, Expr, MetaNameValue};
use syn::parse::Parser;

fn parse_args(attr: proc_macro2::TokenStream) -> Vec<Meta> {
    let parser = syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated;
    parser.parse2(attr).map(|p| p.into_iter().collect()).unwrap_or_default()
}

fn parse_kv_string(metas: &[Meta], key: &str) -> Option<String> {
    for m in metas {
        if let Meta::NameValue(nv) = m {
            if nv.path.is_ident(key) {
                if let Expr::Lit(expr_lit) = &nv.value {
                    if let Lit::Str(s) = &expr_lit.lit {
                        return Some(s.value());
                    }
                }
            }
        }
    }
    None
}

fn parse_kv_expr(metas: &[Meta], key: &str) -> Option<Expr> {
    for m in metas {
        if let Meta::NameValue(nv) = m {
            if nv.path.is_ident(key) {
                return Some(nv.value.clone());
            }
        }
    }
    None
}

fn parse_parameters_json(metas: &[Meta]) -> Option<proc_macro2::TokenStream> {
    for m in metas {
        if let Meta::List(list) = m {
            if list.path.is_ident("parameters_json") {
                return Some(list.tokens.clone());
            }
        }
    }
    None
}

fn parse_annotations(metas: &[Meta]) -> Option<Vec<(String, Expr)>> {
    for m in metas {
        if let Meta::List(list) = m {
            if list.path.is_ident("annotations") {
                let punct = syn::punctuated::Punctuated::<MetaNameValue, syn::Token![,]>::parse_terminated
                    .parse2(list.tokens.clone())
                    .unwrap_or_default();
                let mut pairs = Vec::new();
                for nv in punct {
                    if let Some(ident) = nv.path.get_ident() {
                        let key = ident.to_string();
                        let val = nv.value.clone();
                        pairs.push((key, val));
                    }
                }
                return Some(pairs);
            }
        }
    }
    None
}

#[proc_macro_attribute]
pub fn mcp_tool(attr: TokenStream, item: TokenStream) -> TokenStream {
    let metas = parse_args(attr.into());
    let input_fn = parse_macro_input!(item as ItemFn);

    let fn_name = input_fn.sig.ident.clone();
    let tool_fn_name = format_ident!("{}_tool", fn_name);
    let register_fn_name = format_ident!("{}_register", fn_name);
    let factory_static = format_ident!("__{}_factory", fn_name);

    // Required: name
    let name = parse_kv_string(&metas, "name").unwrap_or_else(|| fn_name.to_string());

    let description = parse_kv_string(&metas, "description");
    let summary = parse_kv_string(&metas, "summary");
    let parameters_expr = parse_kv_expr(&metas, "parameters");
    let parameters_json = parse_parameters_json(&metas);
    let annotations = parse_annotations(&metas);

    // Prepare optional sections
    let description_set = description.as_ref().map(|s| quote! { __def = __def.with_description(#s); });
    let summary_set = summary.as_ref().map(|s| quote! { __def = __def.with_summary(#s); });
    let parameters_set = parameters_expr.as_ref().map(|e| quote! { __def = __def.with_parameters(#e); });
    let parameters_json_set = parameters_json.as_ref().map(|ts| quote! { __def = __def.with_parameters(serde_json::json!({ #ts })); });
    let annotations_set = annotations.as_ref().map(|pairs| {
        let inserts: Vec<proc_macro2::TokenStream> = pairs.iter()
            .map(|(k, v)| quote! { __ann.insert(String::from(#k), #v); })
            .collect();
        quote! {
            let mut __ann = serde_json::Map::<String, serde_json::Value>::new();
            #( #inserts )*
            __def = __def.with_annotations(__ann);
        }
    });

    let output_tokens = quote! {
        #input_fn

        #[allow(non_snake_case)]
        pub fn #tool_fn_name() -> fastmcp_rs::tool::ToolDefinition {
            let mut __def = fastmcp_rs::tool::ToolDefinition::new(#name, #fn_name);
            #description_set
            #summary_set
            #parameters_set
            #parameters_json_set
            #annotations_set
            __def
        }

        #[allow(non_snake_case)]
        pub fn #register_fn_name(server: &fastmcp_rs::FastMcpServer) -> fastmcp_rs::Result<()> {
            server.register_tool(#tool_fn_name())
        }

        #[cfg(feature = "auto-register")]
        #[linkme::distributed_slice(fastmcp_rs::tool::MCP_TOOL_FACTORIES)]
        pub static #factory_static: fastmcp_rs::tool::ToolFactory = || { #tool_fn_name() };
    };

    output_tokens.into()
}